package com.ruoyix.ai.service.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjUtil;
import cn.hutool.core.util.StrUtil;
import com.ruoyix.ai.domain.*;
import com.ruoyix.ai.mapper.AiApiKeyMapper;
import com.ruoyix.ai.mapper.AiChatConversationMapper;
import com.ruoyix.ai.mapper.AiChatModelMapper;
import com.ruoyix.ai.util.AiUtils;
import com.ruoyix.common.core.domain.R;
import com.ruoyix.common.utils.DateUtils;
import com.ruoyix.common.utils.StringUtils;
import com.ruoyix.common.core.page.TableDataInfo;
import com.ruoyix.common.core.domain.PageQuery;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.StreamingChatModel;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.api.ApiUtils;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.stereotype.Service;
import com.ruoyix.ai.domain.bo.AiChatMessageBo;
import com.ruoyix.ai.domain.vo.AiChatMessageVo;
import com.ruoyix.ai.mapper.AiChatMessageMapper;
import com.ruoyix.ai.service.IAiChatMessageService;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.*;

/**
 * AI 聊天消息Service业务层处理
 *
 * @author ruoyi
 * @date 2024-09-21
 */
@RequiredArgsConstructor
@Service
@Slf4j
public class AiChatMessageServiceImpl implements IAiChatMessageService {

    private final AiChatMessageMapper baseMapper;

    private final AiChatConversationMapper aiChatConversationMapper;

    private final AiChatModelMapper aiChatModelMapper;

    private final AiApiKeyMapper aiApiKeyMapper;

    /**
     * 查询AI 聊天消息
     */
    @Override
    public AiChatMessageVo queryById(Long id){
        return baseMapper.selectVoById(id);
    }

    /**
     * 查询AI 聊天消息列表
     */
    @Override
    public TableDataInfo<AiChatMessageVo> queryPageList(AiChatMessageBo bo, PageQuery pageQuery) {
        LambdaQueryWrapper<AiChatMessage> lqw = buildQueryWrapper(bo);
        Page<AiChatMessageVo> result = baseMapper.selectVoPage(pageQuery.build(), lqw);
        return TableDataInfo.build(result);
    }

    /**
     * 查询AI 聊天消息列表
     */
    @Override
    public List<AiChatMessageVo> queryList(AiChatMessageBo bo) {
        LambdaQueryWrapper<AiChatMessage> lqw = buildQueryWrapper(bo);
        return baseMapper.selectVoList(lqw);
    }

    private LambdaQueryWrapper<AiChatMessage> buildQueryWrapper(AiChatMessageBo bo) {
        Map<String, Object> params = bo.getParams();
        LambdaQueryWrapper<AiChatMessage> lqw = Wrappers.lambdaQuery();
        lqw.eq(bo.getConversationId() != null, AiChatMessage::getConversationId, bo.getConversationId());
        lqw.eq(bo.getReplyId() != null, AiChatMessage::getReplyId, bo.getReplyId());
        lqw.eq(bo.getUserId() != null, AiChatMessage::getUserId, bo.getUserId());
        lqw.eq(bo.getRoleId() != null, AiChatMessage::getRoleId, bo.getRoleId());
        lqw.eq(StringUtils.isNotBlank(bo.getType()), AiChatMessage::getType, bo.getType());
        lqw.eq(StringUtils.isNotBlank(bo.getModel()), AiChatMessage::getModel, bo.getModel());
        lqw.eq(bo.getModelId() != null, AiChatMessage::getModelId, bo.getModelId());
        lqw.eq(StringUtils.isNotBlank(bo.getContent()), AiChatMessage::getContent, bo.getContent());
        lqw.eq(bo.getUseContext() != null, AiChatMessage::getUseContext, bo.getUseContext());
        lqw.eq(bo.getDeleted() != null, AiChatMessage::getDeleted, bo.getDeleted());
        return lqw;
    }

    /**
     * 新增AI 聊天消息
     */
    @Override
    public Boolean insertByBo(AiChatMessageBo bo) {
        AiChatMessage add = BeanUtil.toBean(bo, AiChatMessage.class);
        validEntityBeforeSave(add);
        boolean flag = baseMapper.insert(add) > 0;
        if (flag) {
            bo.setId(add.getId());
        }
        return flag;
    }

    /**
     * 修改AI 聊天消息
     */
    @Override
    public Boolean updateByBo(AiChatMessageBo bo) {
        AiChatMessage update = BeanUtil.toBean(bo, AiChatMessage.class);
        validEntityBeforeSave(update);
        return baseMapper.updateById(update) > 0;
    }

    /**
     * 保存前的数据校验
     */
    private void validEntityBeforeSave(AiChatMessage entity){
        //TODO 做一些数据校验,如唯一约束
    }

    /**
     * 批量删除AI 聊天消息
     */
    @Override
    public Boolean deleteWithValidByIds(Collection<Long> ids, Boolean isValid) {
        if(isValid){
            //TODO 做一些业务上的校验,判断是否需要校验
        }
        return baseMapper.deleteBatchIds(ids) > 0;
    }

    public Flux<R<AiChatMessageSendRespVO>> sendChatMessageStream(AiChatMessageSendReqVO sendReqVO, Long userId) {
        // 1.1 校验对话存在
        AiChatConversation conversation = validateConversation(sendReqVO.getConversationId(), userId);
        List<AiChatMessage> historyMessages = baseMapper.selectListByConversationId(conversation.getId());

        // 1.2 校验模型
        AiChatModel model = validateModel(conversation.getModelId());

        // 获取 API Key 并构建流模型
        AiApiKey aiApiKey = aiApiKeyMapper.selectById(model.getKeyId());
        StreamingChatModel chatModel = buildOpenAiChatModel(aiApiKey.getApiKey(), aiApiKey.getUrl());

        // 2. 插入 user 发送消息
        AiChatMessage userMessage = insertUserMessage(conversation, userId, model, sendReqVO);

        // 3. 插入 assistant 接收消息
        AiChatMessage assistantMessage = insertAssistantMessage(conversation, userMessage, model, userId, sendReqVO);

        // 3.2 构建 Prompt，并进行调用
        Prompt prompt = buildPrompt(conversation, historyMessages, model, sendReqVO);
        Flux<ChatResponse> streamResponse = chatModel.stream(prompt);

        // 3.3 流式返回并处理结果
        StringBuilder contentBuilder = new StringBuilder();

        return streamResponse
            .flatMapSequential(chunk -> processChunk(contentBuilder, userMessage, assistantMessage, chunk))
            .doOnComplete(() -> handleStreamCompletion(assistantMessage, contentBuilder))
            .doOnError(throwable -> {
                handleStreamError(assistantMessage, throwable, userId, sendReqVO);
            })
             .onErrorContinue((throwable, obj) -> {
                // 继续流处理并跳过错误元素
                log.error("忽略错误: {}", throwable.getMessage());
            });
    }

    private AiChatConversation validateConversation(Long conversationId, Long userId) {
        AiChatConversation conversation = aiChatConversationMapper.selectById(conversationId);
        if (ObjUtil.notEqual(conversation.getUserId(), userId)) {
            throw new RuntimeException("对话不存在");
        }
        return conversation;
    }

    private AiChatModel validateModel(Long modelId) {
        AiChatModel model = aiChatModelMapper.selectById(modelId);
        if (model == null) {
            throw new RuntimeException("模型不存在");
        }
        return model;
    }

    private AiChatMessage insertUserMessage(AiChatConversation conversation, Long userId, AiChatModel model, AiChatMessageSendReqVO sendReqVO) {
        return createChatMessage(conversation.getId(), null, model, userId, conversation.getRoleId(),
            MessageType.USER, sendReqVO.getContent(), sendReqVO.getUseContext());
    }

    private AiChatMessage insertAssistantMessage(AiChatConversation conversation, AiChatMessage userMessage, AiChatModel model, Long userId, AiChatMessageSendReqVO sendReqVO) {
        return createChatMessage(conversation.getId(), userMessage.getId(), model, userId, conversation.getRoleId(),
            MessageType.ASSISTANT, "", sendReqVO.getUseContext());
    }

    private Mono<R<AiChatMessageSendRespVO>> processChunk(StringBuilder contentBuilder, AiChatMessage userMessage, AiChatMessage assistantMessage, ChatResponse chunk) {
        // 获取 chunk 中的内容
        String newContent = chunk.getResult() != null && !("null".equals(chunk.getResult().getOutput().getContent())) ? chunk.getResult().getOutput().getContent() : "";
        log.info("接收到新的内容块: {}", newContent);

        // 将新内容追加到 StringBuilder 中
        contentBuilder.append(newContent);

        // 构建响应结果
        AiChatMessageSendRespVO aiChatMessageSendRespVO = new AiChatMessageSendRespVO();
        aiChatMessageSendRespVO.setSend(BeanUtil.toBean(userMessage, AiChatMessageSendRespVO.Message.class));

        AiChatMessageSendRespVO.Message message = BeanUtil.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class);
        message.setContent(newContent); // 设置响应内容
        aiChatMessageSendRespVO.setReceive(message);

        // 返回一个 Mono 包含的响应
        return Mono.just(R.ok(aiChatMessageSendRespVO));
    }

    private void handleStreamCompletion(AiChatMessage assistantMessage, StringBuilder contentBuilder) {
        // 输出日志和进行后续操作
        log.info("流处理完成，完整内容为: {}", contentBuilder.toString());

        // 更新数据库中 assistantMessage 的内容
        assistantMessage.setContent(contentBuilder.toString());
        baseMapper.updateById(assistantMessage);
    }

    private void handleStreamError(AiChatMessage assistantMessage, Throwable throwable, Long userId, AiChatMessageSendReqVO sendReqVO) {
        // 处理流中的异常
        log.error("[sendChatMessageStream][userId({}) sendReqVO({}) 发生异常: {}]", userId, sendReqVO, throwable.getMessage());

        // 更新数据库中 assistantMessage 的错误信息
        assistantMessage.setContent(throwable.getMessage());
        baseMapper.updateById(assistantMessage);
    }





    public AiChatMessage sendChatMessageCall(AiChatMessageSendReqVO sendReqVO, Long userId) {
        // 1.1 校验对话存在
        AiChatConversation conversation = aiChatConversationMapper.selectById(sendReqVO.getConversationId());
        if (ObjUtil.notEqual(conversation.getUserId(), userId)) {
            throw new RuntimeException("对话不存在");
        }
        List<AiChatMessage> historyMessages = baseMapper.selectListByConversationId(conversation.getId());
        // 1.2 校验模型
        AiChatModel model = aiChatModelMapper.selectById(conversation.getModelId());

        AiApiKey aiApiKey = aiApiKeyMapper.selectById(model.getKeyId());
        OpenAiChatModel chatModel = buildOpenAiChatModel(aiApiKey.getApiKey(), aiApiKey.getUrl());

        // 2. 插入 user 发送消息
        AiChatMessage userMessage = createChatMessage(conversation.getId(), null, model,
            userId, conversation.getRoleId(), MessageType.USER, sendReqVO.getContent(), sendReqVO.getUseContext());

        // 3.1 插入 assistant 接收消息
        AiChatMessage assistantMessage = createChatMessage(conversation.getId(), userMessage.getId(), model,
            userId, conversation.getRoleId(), MessageType.ASSISTANT, "", sendReqVO.getUseContext());

        // 3.2 构建 Prompt，并进行调用
        Prompt prompt = buildPrompt(conversation, historyMessages, model, sendReqVO);
        ChatResponse call = chatModel.call(prompt);
        String result = call.getResult().getOutput().getContent();

        AiChatMessage aiChatMessage = new AiChatMessage();
        aiChatMessage.setId(assistantMessage.getId());
        aiChatMessage.setContent(result); // 保存完整内容到数据库
        aiChatMessage.setType(MessageType.ASSISTANT.getValue());
        baseMapper.updateById(aiChatMessage);


       return aiChatMessage;
    }



    private static OpenAiChatModel buildOpenAiChatModel(String openAiToken, String url) {
        url = StrUtil.blankToDefault(url, ApiUtils.DEFAULT_BASE_URL);
        OpenAiApi openAiApi = new OpenAiApi(url, openAiToken);
        return new OpenAiChatModel(openAiApi);
    }

    private Prompt buildPrompt(AiChatConversation conversation, List<AiChatMessage> messages,
                               AiChatModel model, AiChatMessageSendReqVO sendReqVO) {
        // 1. 构建 Prompt Message 列表
        List<Message> chatMessages = new ArrayList<>();
        // 1.1 system context 角色设定
        if (StrUtil.isNotBlank(conversation.getSystemMessage())) {
            chatMessages.add(new SystemMessage(conversation.getSystemMessage()));
        }
        // 1.2 history message 历史消息
        List<AiChatMessage> contextMessages = filterContextMessages(messages, conversation, sendReqVO);
        contextMessages.forEach(message -> chatMessages.add(AiUtils.buildMessage(message.getType(), message.getContent())));
        // 1.3 user message 新发送消息
        chatMessages.add(new UserMessage(sendReqVO.getContent()));

        // 2. 构建 ChatOptions 对象
        ChatOptions chatOptions = AiUtils.buildChatOptions(model.getPlatform(), model.getModel(),
            conversation.getTemperature().doubleValue(), conversation.getMaxTokens());
        return new Prompt(chatMessages, chatOptions);
    }

    /**
     * 从历史消息中，获得倒序的 n 组消息作为消息上下文
     *
     * n 组：指的是 user + assistant 形成一组
     *
     * @param messages 消息列表
     * @param conversation 对话
     * @param sendReqVO 发送请求
     * @return 消息上下文
     */
    private List<AiChatMessage> filterContextMessages(List<AiChatMessage> messages,
                                                        AiChatConversation conversation,
                                                      AiChatMessageSendReqVO sendReqVO) {
        if (conversation.getMaxContexts() == null || ObjUtil.notEqual(sendReqVO.getUseContext(), Boolean.TRUE)) {
            return Collections.emptyList();
        }
        List<AiChatMessage> contextMessages = new ArrayList<>(conversation.getMaxContexts().intValue() * 2);
        for (int i = messages.size() - 1; i >= 0; i--) {
            AiChatMessage assistantMessage = CollUtil.get(messages, i);
            if (assistantMessage == null || assistantMessage.getReplyId() == null) {
                continue;
            }
            AiChatMessage userMessage = CollUtil.get(messages, i - 1);
            if (userMessage == null || ObjUtil.notEqual(assistantMessage.getReplyId(), userMessage.getId())
                || StrUtil.isEmpty(assistantMessage.getContent())) {
                continue;
            }
            // 由于后续要 reverse 反转，所以先添加 assistantMessage
            contextMessages.add(assistantMessage);
            contextMessages.add(userMessage);
            // 超过最大上下文，结束
            if (contextMessages.size() >= conversation.getMaxContexts() * 2) {
                break;
            }
        }
        Collections.reverse(contextMessages);
        return contextMessages;
    }

    private AiChatMessage createChatMessage(Long conversationId, Long replyId,
                                            AiChatModel model, Long userId, Long roleId,
                                            MessageType messageType, String content, Boolean useContext) {


        AiChatMessage message = new AiChatMessage();
        message.setConversationId(conversationId);
        message.setReplyId(replyId);
        message.setModel(model.getModel());
        message.setModelId(model.getId());
        message.setUserId(userId);
        message.setRoleId(roleId);
        message.setType(messageType.getValue());
        message.setContent(content);
        message.setUseContext(useContext ? 1 : 0);
        message.setCreateTime(DateUtils.getNowDate());
        baseMapper.insert(message);
        return message;
    }
}
